#!/usr/bin/env python3

import argparse
from glob import glob
import os
import pendulum
from jinja2 import Template
import re
import sys
import config_wrfda
sys.path.append(f'{os.path.dirname(os.path.realpath(__file__))}/../utils')
from utils import cli, check_files, search_files, run, submit_job, parse_config, get_value, has_key

scripts_root = os.path.dirname(os.path.realpath(__file__))

def run_wrfda_3dvar(work_root, wrfda_root, config, args, wrf_work_dir=None, force=False, dom='d01', tag=None, fg=None):
	start_time = config['custom']['start_time']
	datetime_fmt = 'YYYY-MM-DD_HH:mm:ss'
	start_time_str = start_time.format(datetime_fmt)
	max_dom = config['domains']['max_dom']

	if not wrf_work_dir:
		if tag != None:
			wrf_work_dir = f'{work_root}/{tag}/wrf'
		else:
			wrf_work_dir = f'{work_root}/wrf'

	if tag != None:
		obsproc_work_dir = f'{work_root}/{tag}/wrfda/obsproc'
	else:
		obsproc_work_dir = f'{work_root}/wrfda/obsproc'

	dom_str = dom
	if tag != None:
		wrfda_work_dir = f'{work_root}/{tag}/wrfda/{dom_str}'
	else:
		wrfda_work_dir = f'{work_root}/wrfda/{dom_str}'
	if not os.path.isdir(wrfda_work_dir): os.mkdir(wrfda_work_dir)
	old_dir = os.getcwd()
	os.chdir(wrfda_work_dir)

	cli.stage(f'Run da_wrfvar.exe at {wrfda_work_dir} ...')

	if os.path.isfile(f'wrfvar_output_{start_time_str}') and not args.force and not force:
		run(f'ls -l wrfvar_output_{start_time_str}')
		cli.notice(f'wrfvar_output_{start_time_str} already exist.')
		return

	run(f'ln -sf {wrfda_root}/run/LANDUSE.TBL {wrfda_work_dir}')

	if not os.path.isfile('namelist.input'):
		cli.error('namelist.input has not been generated! Run config_wrfda.py.')

	# BE matrix
	run(f'rm -f be.dat')
	if 'cv_options' in config['wrfvar7']:
		if args.be_data:
			run(f'ln -sf {args.be_data} be.dat')
		else:
			be_work_dir = os.path.abspath(work_root) + '/be/' + dom_str
			if not os.path.isdir(be_work_dir):
				be_work_dir = os.path.abspath(work_root) + '/../be/' + dom_str

			if config['wrfvar7']['cv_options'] == 5:
				if not os.path.isfile(f'{be_work_dir}/be.dat.cv5'):
					cli.error(f'BE matrix {be_work_dir}/be.dat.cv5 does not exist!')
				run(f'ln -sf {be_work_dir}/be.dat.cv5 be.dat')
			elif config['wrfvar7']['cv_options'] == 6:
				if not os.path.isfile(f'{be_work_dir}/be.dat.cv6'):
					cli.error(f'BE matrix {be_work_dir}/be.dat.cv6 does not exist!')
				run(f'ln -sf {be_work_dir}/be.dat.cv6 be.dat')
			elif config['wrfvar7']['cv_options'] == 7:
				if not os.path.isfile(f'{be_work_dir}/be.dat.cv7'):
					cli.error(f'BE matrix {be_work_dir}/be.dat.cv7 does not exist!')
				run(f'ln -sf {be_work_dir}/be.dat.cv7 be.dat')
	if not os.path.exists('./be.dat'):
		run(f'ln -sf {wrfda_root}/var/run/be.dat.cv3 be.dat')

	# First guess
	# TODO: Assume there is only one domain to be assimilated.
	fg_path = f'{wrfda_work_dir}/fg'
	if fg != None:
		run(f'cp {fg} {fg_path}')
	elif not os.path.isfile(fg_path):
		expected_files = ['{}/wrfinput_d{:02d}_{}'.format(wrf_work_dir, i + 1, start_time_str) for i in range(max_dom)]
		if not check_files(expected_files):
			cli.error('real.exe or da_update_bc.exe wasn\'t executed successfully!')
		run(f'cp {wrf_work_dir}/wrfinput_{dom_str}_{start_time_str} {fg_path}')

	# Observation data
	if get_value(config, ('custom', 'wrfda', 'type'), '3dvar') == '3dvar':
		# Radar observation:
		if 'use_radarobs' in config['wrfvar4'] and config['wrfvar4']['use_radarobs']:
			radar_path = f'{obsproc_work_dir}/obs.radar.{start_time.format("YYYYMMDDHHmm")}'
			if os.path.isfile(radar_path):
				run(f'ln -sf {radar_path} ob.radar')
			else:
				cli.error(f'There is no Radar file {radar_path}!')
		# Conventional observation:
		if config['wrfvar3']['ob_format'] == 2 and os.path.isfile(f'{obsproc_work_dir}/obs_gts_{start_time.format(datetime_fmt)}.3DVAR'):
			# LITTLE_R conventional data
			run(f'ln -sf {obsproc_work_dir}/obs_gts_{start_time.format(datetime_fmt)}.3DVAR ob.ascii')
		elif config['wrfvar3']['ob_format'] == 1 and config['custom']['wrfda']['prepbufr_source'] == 'gdas':
			if not args.prepbufr_root:
				cli.error(f'Option --prepbufr-root should be provided!')
			# PREPBUFR conventional data
			gdas_file_path = f'{args.prepbufr_root}/gdas.{start_time.format("YYYYMMDD")}/gdas.t{start_time.hour:02}z.prepbufr.nr'
			if not os.path.isfile(gdas_file_path):
				cli.error(f'{gdas_file_path} does not exist!')
			run(f'ln -sf {gdas_file_path} ob.bufr')
	else:
		cli.error('Call run_wrfda_3dvar, so set config->wrfda->type to 3dvar!')

	if os.path.isfile(f'{wrfda_work_dir}/wrfvar_output_{start_time_str}') and not args.force:
		cli.notice(f'{wrfda_work_dir}/wrfvar_output_{start_time_str} already exists.')
		return

	submit_job(f'{wrfda_root}/var/build/da_wrfvar.exe', args.np, config, args, wait=True)

	expected_files = [f'wrfvar_output', 'statistics']
	if not check_files(expected_files):
		# Check if the failure is caused by parallel computing? Such as cv_options is zero in some process.
		if search_files('rsl.error.*', 'Invalid CV option chosen:  cv_options =    0'):
			cli.warning('Failed to run da_wrfvar.exe in parallel. Try to run in serial.')
			submit_job(f'{wrfda_root}/var/build/da_wrfvar.exe', 1, config, args, wait=True)
			if not check_files(expected_files):
				cli.error(f'Still failed! See {wrfda_work_dir}/rsl.error.0000.')
		else:
			cli.error(f'Failed! See {wrfda_work_dir}/rsl.error.0000.')
	else:
		print(open('statistics').read())
		# run(f'ncl -Q {scripts_root}/../plots/plot_cost_grad_fn.ncl')
		run(f'cp wrfvar_output wrfvar_output_{start_time_str}')
		cli.notice('Succeeded.')
	os.chdir(old_dir)
